import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms as transform
from torch.utils.data import DataLoader, Dataset
import numpy as np
import matplotlib.pyplot as plt
import tqdm

import sys
import os

# 获取父目录的路径（project/）并添加到搜索路径
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(parent_dir)

from common.utils import seed_everything
from denoising_config import *
from denoising_data import *
from denoising_model import *
from denoising_engine import train_step, test_step

if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    #指定随机数种子
    seed_everything(SEED)

    transform = transform.Compose([
    transform.Resize((IMG_HEIGHT,IMG_WIDTH)),
    transform.ToTensor()])

    #加载数据集
    dataset = ImageDataset(image_dir=IMG_PATH, transform=transform)
    train_dataset, test_dataset = torch.utils.data.random_split(dataset, [TRAIN_SPLIT, TEST_SPLIT])
    #创建数据加载器
    train_loader =DataLoader(dataset=train_dataset,
                             batch_size=TRAIN_BATCH_SIZE,shuffle=True,drop_last=True)
    test_loader = DataLoader(dataset=test_dataset,
                            batch_size=TEST_BATCH_SIZE,shuffle=False)
    #训练模型
    denoiser = ConvDenoiser().to(device)
    loss = nn.MSELoss()
    optimizer = optim.Adam(denoiser.parameters(), lr=LEARNING_RATE)
    print("开始训练...")
    min_test_loss = 999
    for epoch in tqdm.tqdm(range(EPOCHS)):
        train_loss = train_step(model=denoiser,
                                data_loader=train_loader,
                                loss_fn=loss,
                                optimizer=optimizer,
                                device=device)
        test_loss = test_step(model=denoiser,
                              test_loader=test_loader,
                              loss_fn=loss,
                              device=device)
        print(f'Epoch {epoch+1}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}')

        #保存测试集上表现最好的模型
        if test_loss < min_test_loss:
            min_test_loss = test_loss
            print("保存模型...")
            torch.save(denoiser.state_dict(), DENOISER_MODEL_FILE)
    print("训练完成！")
    
    